import torch
import torch.nn as nn
from torchvision import transforms, datasets
import json
import matplotlib.pyplot as plt
import os
import torch.optim as optim
# from model import resnet34, resnet101
import torchvision.models.resnet

# 模块类
class BasicBlock(nn.Module):

    expansion = 1

    def __init__(self, in_channel, out_channel, stride=1, downsample=None):  # downsample对应虚线残差结构
        super(BasicBlock, self).__init__()
        # 特征提取，特征变换,调整通道数
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel,
                               kernel_size=3, stride=stride, padding=1,bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)  # BN处理
        self.relu = nn.ReLU()

        # 特征提取，特征转换
        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
                               kernel_size=3, stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.downsample = downsample

    def forward(self, x):
        identity = x  # 存放捷径上的输出值
        if self.downsample is not None:
            identity = self.downsample(x)

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out+=identity

        return out

class Bottleneck(nn.Module):

    # 模块中最后一层卷积层卷积核的倍数
    # expansion = 4  # 每个conv的卷积和个数的倍数

    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        super(Bottleneck, self).__init__()

        # 这一步的目的是什么
        # self.conv1 = nn.Conv2d(in_channel=in_channel, out_channel=out_channel,
        #                        kernel_size=1, stride=1, padding=1,bias=False)
        # self.bn1 = nn.BatchNorm2d(out_channel)  # BN处理
        # self.relu = nn.ReLU(inplace=True)

        # 特征提取，特征转换
        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
                               kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU(inplace=True)

        self.conv3 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel,
                               kernel_size=1, stride=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU(inplace=True)

        self.downsample = downsample

    def forward(self, x):
        identity = x  # 存放捷径上的输出值
        if self.downsample is not None:
            identity = self.downsample(x)

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out+=identity

        return out

class ResNet(nn.Module):

    def __init__(self, block, block_num, num_classes=1000):
        '''
        :param block:模块类
        :param block_num: 一个列表，指定了每个层级（layer1,layer2,layer3,layer4）中残差快的数量中
        :param num_classes:用于全连接层
        :include_top:指定是否包含全连接层
        '''

        super(ResNet, self).__init__()

        # 模型输入通道
        self.in_channel = 64

        # 调整通道数，特征提取，特征转换
        self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2,
                               padding=3, bias=False)  # 填充3个单位
        self.bn1 = nn.BatchNorm2d(self.in_channel)
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  # 最大池化层，减小空间维度和参数个数，防止过拟合

        # 定义主要层级，每个层级由多个残差块组成
        self.layer1 = self._make_layer(block, 64, block_num[0])
        self.layer2 = self._make_layer(block, 128, block_num[1], stride=2)
        self.layer3 = self._make_layer(block, 256, block_num[2], stride=2)
        self.layer4 = self._make_layer(block, 512, block_num[3], stride=2)

        # 如果include_top=True添加一个全局平均池化层和一个全连接层

        self.avgpool =  nn.AdaptiveAvgPool2d((1, 1))  # output size = (1, 1)
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        # 初始化函数，考虑了前向传播和反向传播时激活值的方差和梯度的方差，旨在保持信号在网络中的流动，避免梯度消失或爆炸问题
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

    def _make_layer(self, block, channel, block_num, stride=1):
        '''
        :param block:
        :param channel: 当前层级的输出通道数
        :param block_num: 当前层级中残差快的数量
        :param stride:
        :return:
        '''

        downsample = None

        # 为什么满足这两种情况要添加残差快
        # 如果stride不为1或者当前输入通道数与期望输出通道数不匹配就添加一个下采样
        # 没达到预期标准或为深层次的卷积层
        if stride != 1 or self.in_channel != channel * block.expansion:
            downsample = nn.Sequential(nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
                                       nn.BatchNorm2d(channel * block.expansion))

        layers = []
        layers.append(block(self.in_channel, channel * block.expansion, downsample=downsample, stride=stride))
        self.in_channel = channel * block.expansion

        for _ in range(1, block_num):
            layers.append(block(self.in_channel, channel))
            self.in_channel = channel

        return nn.Sequential(*layers)  #  一个解包操作符，将layer列表（或元组等可迭代对象）中的元素解包成位置参数，然后传给nn.Sequential的构造函数。这样，nn.Sequential接受多个层作为输入

    def forward(self, x):
        # 调整通道数，特征提取，特征转换
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        # 主要层次
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        if self.include_top:
            x = self.avgpool
            x = torch.flatten(x, 1)
            x = self.fc(x)

        return x

def resnet34(num_classes=1000):
    return ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes)

def resnet101(num_classes=1000):
    return ResNet(BasicBlock, [3, 4, 23, 3], num_classes=num_classes)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

data_transform = {
    "train":transforms.Compose([# transforms.RandomResizedCrop(224),
                               # transforms.RandomHorrizontalFlip(),
                                transforms.Resize((3324, 3324)),
                               transforms.ToTensor(),
                               transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),  # 来自网格参数
    "val":transforms.Compose([# transforms.Resize(256),  # 将最小边长缩放到256
                               # transforms.CenterCrop(224),
                                transforms.Resize((3324, 3324)),
                               transforms.ToTensor(),
                               transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),  # 来自网格参数
}

img_path = "D:\\APP_Detection\\image_data\\dex_image\\"

train_dataset = datasets.ImageFolder(root=img_path + "dex_train_image", transform=data_transform["train"])
train_num = len(train_dataset)

# 转换类别索引为类别名称字典
class_to_idx = train_dataset.class_to_idx
idx_to_class = dict((val, key) for key, val in class_to_idx.items())

# 将索引写入json文件
json_str = json.dumps(idx_to_class, indent=4)
with open('class_indices.json', 'w') as json_file:
    json_file.write(json_str)

batch_size = 5
train_loader = torch.utils.data.DataLoader(train_dataset,
                                           batch_size=batch_size, shuffle=True,
                                           num_workers=0)

# 验证集
validate_dataset = datasets.ImageFolder(root=img_path + "dex_train_image",
                                        transform=data_transform["val"])
val_num = len(validate_dataset)
validate_loader = torch.utils.data.DataLoader(validate_dataset,
                                              batch_size=batch_size, shuffle=False,
                                              num_workers=0)
# net = resnet34()
net = resnet34(num_classes=5)

# load pretrain weights

# model_weight_path = "./resnet34-pre.pth"
# missing_keys, unexpected_keys = net.load_state_dict(torch.load(model_weight_path), strict=False)#载入模型参数

# for param in net.parameters():
#     param.requires_grad = False
# change fc layer structure

# inchannel = net.fc.in_features
# net.fc = nn.Linear(inchannel, 5)


net.to(device)

loss_function = nn.CrossEntropyLoss()  # 用于多分类问题的损失函数
optimizer = optim.Adam(net.parameters(), lr=0.0001)

best_acc = 0.0  # 用于记录模型在验证集上的最佳准确率
save_path = "D:\\resnet\\resnet34_best"  # 保存模型权重的路径
for epoch in range(3):  # 每个epoch代表整个训练集被遍历一次
    # train
    net.train()
    running_loss = 0.0
    for step, data in enumerate(train_loader, start=0):
        images, labels = data
        optimizer.zero_grad()
        logits = net(images.to(device))
        loss = loss_function(logits, labels.to(device))
        loss.backward()
        optimizer.step()  # 更新模型参数

        # print statistics
        running_loss += loss.item()
        # print train process
        rate = (step + 1) / len(train_loader)
        a = "*" * int(rate * 50)
        b = "." * int((1 - rate) * 50)
        print("\rtrain loss: {:^3.0f}%[{}->{}]{:.4f}".format(int(rate * 100), a, b, loss), end="")
    print()

    # validate
    net.eval()  # 将模型设置为评估模式
    acc = 0.0  # 初始化准确率累加器
    with torch.no_grad():
        for val_data in validate_loader:
            val_images, val_labels = val_data
            outputs = net(val_images.to(device))  # eval model only have last output layer
            # loss = loss_function(outputs, test_labels)
            predict_y = torch.max(outputs, dim=1)[1]
            acc += (predict_y == val_labels.to(device)).sum().item()
        val_accurate = acc / val_num
        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)
        print('[epoch %d] train_loss: %.3f  test_accuracy: %.3f' %
              (epoch + 1, running_loss / step, val_accurate))

print('Finished Training')

#预测
